from neutron import Plugin, User
from fastapi import WebSocket
from starlette import websockets

import base64

from . import term
from server import config

app = Plugin(__file__)

term.init()


@app.route.get("/new", level=2)
async def get_terminal_id():
    # 获取可以绑定的termid
    term.bash_id += 1
    return {"termid": term.bash_id}


@app.route.ws("/ws")
async def load_terminal(user: User, websocket: WebSocket, termid: int):
    # websocket链接到termid
    try:
        await websocket.accept()

        if (user.demo):
            await websocket.send_json({"type": "data", "data": base64.b64encode("\n\033[0m\033[47;30m * \033[44;30m 演示模式下禁用终端 \033[0m\n\n".encode("utf-8")).decode()})
            await websocket.close()

        # 设置输出回调绑定到当前ws
        async def oncall(text_b):
            await websocket.send_json({"type": "data", "data": base64.b64encode(text_b).decode()})
        term.termcallbacks[termid] = oncall

        if (termid not in term.bash_lst):
            # 创建终端会话
            term.termoutput[termid] = b""

            # 设置内容回调
            async def callbacks(text_b):
                if (termid not in term.termoutput):
                    return
                term.termoutput[termid] += text_b
                if (len(term.termoutput[termid]) > 3000):
                    term.termoutput[termid] = term.termoutput[termid][-3000:]
                await term.termcallbacks[termid](text_b)
            await term.run_bash(callbacks, termid)
        else:
            # 还原历史记录
            await websocket.send_json({"type": "data", "data": base64.b64encode(term.termoutput[termid]+"\n\033[0m\033[47;30m * \033[44;30m 还原的历史 \033[0m\n\n".encode("utf-8")).decode()})

        while True:
            # 接受数据
            data = await websocket.receive_json()
            if (data["type"] == "data"):  # 命令
                await term.write_bash(termid, base64.b64decode(data["data"]))
            elif (data["type"] == "resize"):  # 缩放
                await term.set_term_size(termid, data["data"][0], data["data"][1])
            elif (data["type"] == "close"):  # 关闭
                await term.close_bash(termid)
                break
    except (RuntimeError, websockets.WebSocketDisconnect):
        # 断连
        if (termid):
            term.termcallbacks[termid] = term.empty
